#pragma once

#include "AssocState.h"

namespace AsyncLib
{

template<typename T>
class CAwaiter
{
public:
	CAwaiter(std::shared_ptr<CAssocState<T>> pState) : m_pState(std::move(pState))
	{
		assert(m_pState);
	}

	CAwaiter(const CAwaiter&) = delete;
	CAwaiter& operator=(const CAwaiter&) = delete;

	bool await_ready() const
	{
		return m_pState->IsCompleted();
	}

	bool await_suspend(std::coroutine_handle<> handle) const
	{
		auto status = m_pState->SetCallback([handle](bool bResult)
		{
			if (bResult)
				handle.resume();
			else
				handle.destroy();
		});

		if (status == CAssocStateBase::Canceled)
			handle.destroy();

		// Continue to await_resume() without suspending when Completed
		return status != CAssocStateBase::Completed;
	}

	decltype(auto) await_resume() const
	{
		return m_pState->GetValue();
	}

private:
	std::shared_ptr<CAssocState<T>> m_pState;
};

} // namespace AsyncLib
